"""
破解 zip 包的解压密码
使用暴力破解, 速度很慢
@date 2019/10/27
@author hujing
"""
import zipfile
import threading
import string
from progressbar import *


class CheckZipPassword:
    def __init__(self, zip_path: str, pwd_min_len: int = 1, pwd_max_len: int = 10,
                 pwd_word_list: list = string.digits + string.ascii_letters + string.punctuation,
                 thread_num: int = 10, check_path='/tmp'):
        """
        :param zip_path: 压缩包路径
        :param pwd_min_len: 检查密码的最短长度
        :param pwd_max_len: 检查密码的最大长度
        :param pwd_word_list: 进行排列组合的字符数组
        :param thread_num: 开启线程数量
        :param check_path: 测试密码时解压的临时路径
        """
        self.zip_file = zipfile.ZipFile(zip_path)
        self.zip_password = None  # 做种找到的密码
        self.pwd_min_len = pwd_min_len  # 密码遍历的开始长度
        self.pwd_max_len = pwd_max_len  # 密码遍历的终止长度
        self.pwd_word_list = pwd_word_list  # 密码遍历时组合包含的字符
        self.thread_num = thread_num
        self.check_path = check_path
        self.pwd_iter = self.foreach_check_password()
        self.iter_lock = threading.Lock()
        pass

    def run(self):
        """
        启动破解线程, 返回破解后的密码
        :return:
        """
        thread_list = []
        # 启动线程
        for i in range(self.thread_num):
            item_t = threading.Thread(target=self.start_thread, args=(i,))
            item_t.start()
            thread_list.append(item_t)
        # 等待线程结束
        for item_thread in thread_list:
            item_thread.join()

        return self.zip_password

    def start_thread(self, id):
        """
        进行密码检查的线程
        :return:
        """

        while True:
            # 判断是否已经找到了最终密码
            if self.zip_password:
                break
            # 拿一个密码
            try:
                with self.iter_lock:
                    check_pwd = next(self.pwd_iter)
            except StopIteration:
                break
            # 检查密码是否正确
            if self.check_password(check_pwd, self.check_path + '/' + str(id)):
                self.zip_password = check_pwd

    def foreach_check_password(self):
        """
        获取待测试的密码. 线程不安全, 使用时请自行加锁
        :return:
        """
        # 总数量
        total_num = 0
        for item_len in range(self.pwd_min_len, self.pwd_max_len + 1):
            total_num += len(self.pwd_word_list) ** item_len
        # 进度条
        bar_widgets = [
            f'总数: {total_num} ',
            Bar(left='[', right=']'),
            Percentage(),
            " ",
            Timer(),
            " ",
            ETA(),
        ]
        bar = ProgressBar(widgets=bar_widgets).start()

        yield_num = 0  # 已经返回的数量
        # 依次检查每个密码长度
        for item_len in range(self.pwd_min_len, self.pwd_max_len + 1):
            for item in self.foreach_all_group(item_len):
                yield_num += 1
                bar.update(int((yield_num / total_num) * 100))
                # print(f'{yield_num}/{total_num}', end="\r")
                yield item

    def foreach_all_group(self, pwd_len: int):
        """
        获取指定长度的所有密码组合
        :param pwd_len:
        :return:
        """
        if pwd_len == 1:
            for item in self.pwd_word_list:
                yield item
        elif pwd_len <= 0:
            return
        else:
            for item_word in self.pwd_word_list:
                for item in self.foreach_all_group(pwd_len - 1):
                    yield item_word + item

    def check_password(self, password, check_path) -> bool:
        """
        测试 zip 密码是否正确
        :return:
        """
        try:
            self.zip_file.extractall(path=check_path, pwd=password.encode('utf-8'))
            return True
        except:
            return False


if __name__ == '__main__':
    zip_path = './b.zip'
    zip_pwd = CheckZipPassword(
        zip_path=zip_path,
        pwd_min_len=6,
        pwd_max_len=10,
        thread_num=500,
    ).run()
    print("\n")
    print(f'查询到结果: {zip_pwd}')
